library(dplyr)
library(tidyr)
library(stringr)
library(ComplexHeatmap)
library(circlize)

if (!file.exists("output")) {
  dir.create("output")
}

# load data

source("HNC_metadata_tidy.R")

sbs_burden <- read.csv("Input_mutational_burden/Input_SBS_mutational_burden.csv")

# Likely driver mutations and CN events identified in HNC. 
HNCdrivers <- read.csv("input_drivers/Input_HNC_driver_mutations.csv") %>% select(-X)
HNCdrivers$Effect <- factor(HNCdrivers$Effect, labels = c("ESS splice", "frameshift", "inframe", "missense", "nonsense", "start lost"))
HNCdrivers$Effect <- factor(HNCdrivers$Effect, levels = c("missense", "nonsense", "frameshift", "ESS splice", "inframe", "start lost"))

CNdrivers <- read.csv("input_drivers/Input_HNC_CN_drivers.csv") %>% rename("donor_id" = "X")
data <- data %>% filter(donor_id %in% CNdrivers$donor_id)

CNevents <- NULL
for (g in 2:ncol(CNdrivers)) {
  cn <- data.frame(Gene = names(CNdrivers)[g], CNdrivers[!is.na(CNdrivers[g]), c(1, g)])
  names(cn) <- c("Gene", "Sample", "Effect")
  CNevents <- rbind(CNevents, cn)
}

# Select all alterations (CN and mutationss) found in CN-altered genes
ALLevents <- HNCdrivers %>%
  filter(Gene %in% names(CNdrivers)[-1], Sample %in% CNdrivers$donor_id) %>%
  select(Gene, Sample, Effect) %>%
  distinct() %>%
  rbind(CNevents)

# Generate a data frame with altered cancer genes per sample
myMut <- ALLevents %>%
  select(Sample, Gene) %>%
  distinct()
muts <- data["donor_id"]
for (g in unique(ALLevents$Gene)) {
  x <- myMut %>%
    filter(Gene == g) %>%
    right_join(data %>% select(donor_id), by = c("Sample" = "donor_id")) %>%
    mutate(Gene = as.character(Gene)) %>%
    mutate(Gene = ifelse(Gene == g, 1, 0))
  x[is.na(x)] <- 0
  colnames(x)[2] <- g
  muts <- left_join(muts, x, by = c("donor_id" = "Sample"))
}

data <- data %>%
  left_join(sbs_burden, by = "donor_id") %>%
  left_join(muts, by = "donor_id") %>%
  mutate_at(vars(names(muts)[-1]), factor) %>%
  mutate(alcohol_ever = factor(alcohol_ever, labels = c("No", "Yes")))

# Calculate percentages of altered samples
percalt <- data.frame()
for (d in unique(ALLevents$Gene)) {
  x <- data %>%
    group_by(get(d)) %>%
    tally() %>%
    mutate(perc = prop.table(n) * 100, .after = "n")
  mut <- data.frame(gene = d, n = x$n[2], perc = round(x$perc[2], 2))
  percalt <- rbind(percalt, mut)
}

# Select frequently-altered cancer genes (e.g. >2% of samples)
drivers <- percalt[percalt$perc > 2, "gene"]

# Generate matrices by mutation effect
MutList <- list()
for (vtype in unique(ALLevents$Effect)) {
  myMut <- ALLevents %>%
    filter(Effect == vtype) %>%
    select(Sample, Gene) %>%
    distinct()
  muts <- data["donor_id"]
  for (g in 1:length(drivers)) {
    x <- myMut %>%
      filter(Gene == drivers[g]) %>%
      right_join(data %>% select(donor_id), by = c("Sample" = "donor_id")) %>%
      mutate(Gene = as.character(Gene)) %>%
      mutate(Gene = ifelse(Gene == drivers[g], 1, 0))
    x[is.na(x)] <- 0
    colnames(x)[2] <- drivers[g]
    muts <- left_join(muts, x, by = c("donor_id" = "Sample"))
  }
  MutList[[vtype]] <- t(as.matrix(muts %>% tibble::column_to_rownames("donor_id")))
}

MutList <- unify_mat_list(MutList, default = 0)

#  Annotations
col <- c(`ESS splice` = "#E69F00", frameshift = "#56B4E9", nonsense = "#D55E00", missense = "#009E73", inframe = "#0072B2", `start lost` = "#F0E442", Deleted = "#005AB5", Amplified = "#DC3220")

top_anno <- HeatmapAnnotation(
  "Mutation \n burden" = anno_barplot(data$Mutational.burden, baseline = 0),
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 10),
  border = T,
  height = unit(1, "cm")
)

bottom_anno <- HeatmapAnnotation(
  "Tobacco" = as.vector(data$tobacco_ever),
  "Alcohol" = as.vector(data$alcohol_ever),
  "HPV" = as.vector(data$hpv_pos_opc),
  "Subsite" = as.vector(data$subsite),
  show_legend = c(
    "Tobacco" = T,
    "Alcohol" = T,
    "HPV" = T,
    "Subsite" = T
  ),
  col = list(
    "Tobacco" = c("No" = "#ffffcc", "Yes" = "#225ea8"),
    "Alcohol" = c("No" = "#ffffcc", "Yes" = "#225ea8"),
    "HPV" = c("Negative" = "#ffffcc", "Positive" = "#225ea8"),
    "Subsite" = c("Larynx" = "#225ea8", "Hypopharynx" = "#41b6c4", "OPC" = "#a1dab4", "OC" = "#ffffcc")
  ),
  gap = unit(c(0, 0, 1), "mm"),
  border = T, na_col = "#e6e6e6",
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 10),
  simple_anno_size = unit(0.4, "cm")
)

# Oncoplot
HM <- oncoPrint(MutList,
  alter_fun_is_vectorized = T,
  alter_fun = list(
    `ESS splice` = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["ESS splice"], col = NA)),
    frameshift = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["frameshift"], col = NA)),
    nonsense = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["nonsense"], col = NA)),
    missense = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["missense"], col = NA)),
    inframe = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["inframe"], col = NA)),
    `start lost` = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.9, gp = gpar(fill = col["start lost"], col = NA)),
    Deleted = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.4, gp = gpar(fill = col["Deleted"], col = NA)),
    Amplified = function(x, y, w, h) grid.rect(x, y, w * 0.9, h * 0.4, gp = gpar(fill = col["Amplified"], col = NA))
  ),
  top_annotation = top_anno,
  bottom_annotation = bottom_anno,
  column_split = factor(data$tob_alc),
  column_gap = unit(1, "mm"),
  col = col,
  width = unit(18, "cm"),
  height = unit(4, "cm"),
  row_names_side = "left", pct_side = "right",
  pct_digits = 1,
  border = T,
  column_title = c("HNC (n = 242)"),
  row_names_gp = gpar(fontsize = 8), pct_gp = gpar(fontsize = 8),
)

pdf(
  file = paste0("ExtendedDataFigure_6b_", Sys.Date(), ".pdf", sep = ""),
  width = 10, height = 15
)
draw(HM,
  heatmap_legend_side = "bottom",
  annotation_legend_side = "bottom",
  merge_legend = TRUE,
)
dev.off()

# Fisher test
factors <- c("tob_alc")

fisher_results <- c()
for (f in factors) {
  for (s in drivers) {
    t <- table(data[, s], data[, f])
    if (ncol(t) < 2) {
      print(paste("Variable", f, "for driver", s, "has been removed due to <2 categories"))
    } else {
      test <- tryCatch(
        {
          fisher.test(t)
        },
        error = function(e) {
          NULL
        }
      )

      if (is.null(test)) {
        print(paste0("Variable", f, "has been removed due to ERROR in fisher test"))
        next
      }

      result <- data.frame(
        driver = s,
        factor = f,
        fisher_p = test$p.value
      )
      fisher_results <- rbind(fisher_results, result)
    }
  }
}

fisher_sig <- fisher_results %>%
  group_by(factor) %>%
  mutate(
    driver = factor(driver, levels = drivers),
    q_val = p.adjust(fisher_p, "fdr")
  ) %>%
  ungroup() %>%
  arrange(driver)
